# -*- coding: utf-8 -*-
"""
Webman - 请求记录与响应回放系统（纯 Text 存储版）
所有字段使用 Text 存储 JSON 字符串，兼容性更强
"""

from flask import Flask, request, jsonify
from sqlalchemy import (
    create_engine, Column, Integer, String, Text, DateTime, Index, exc
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import logging
from logging.handlers import RotatingFileHandler
import json

# ==================== 配置 ====================
DB_USER = 'root'
DB_PASS = 'huanghai2010'  # 你的 MySQL 密码
DB_HOST = 'localhost'
DB_PORT = 3306
DB_NAME = 'webman'

SERVER_PORT = 9000

LOG_FILE = 'requests.log'
LOG_MAX_SIZE = 50 * 1024 * 1024  # 50MB
LOG_BACKUP_COUNT = 5

DATABASE_URI = (
    f"mysql+pymysql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
    "?charset=utf8mb4&autocommit=true"
)

# ==================== 初始化 Flask ====================
app = Flask(__name__)

# 配置日志
handler = RotatingFileHandler(LOG_FILE, maxBytes=LOG_MAX_SIZE, backupCount=LOG_BACKUP_COUNT)
handler.setFormatter(logging.Formatter(
    '%(asctime)s [%(levelname)s] %(message)s'
))
app.logger.addHandler(handler)
app.logger.setLevel(logging.INFO)
logger = app.logger

# ==================== 数据库初始化 ====================
engine = None
SessionLocal = None
Base = declarative_base()

# ✅ 定义模型（全部结构化字段使用 Text + JSON 序列化）
class WebRequest(Base):
    __tablename__ = 'web'

    id = Column(Integer, primary_key=True, autoincrement=True)
    path = Column(String(500), nullable=False)
    method = Column(String(10), nullable=False)
    query_string = Column(Text)  # URL 查询参数原始字符串
    header = Column(Text)       # JSON 字符串
    args = Column(Text)         # JSON 字符串 (request.args)
    form = Column(Text)         # JSON 字符串 (request.form)
    json = Column(Text)         # JSON 字符串 (request.get_json())
    response = Column(Text, default='{"message": "welcome!"}')  # JSON 字符串
    created_at = Column(DateTime, default=datetime.now)

    __table_args__ = (
        Index('idx_method_path', 'method', 'path', unique=True),
    )

def init_database():
    global engine, SessionLocal
    try:
        engine = create_engine(DATABASE_URI, pool_pre_ping=True, echo=False)
        SessionLocal = sessionmaker(bind=engine)
        Base.metadata.create_all(engine)
        logger.info("✅ 数据库连接成功，表已创建或确认存在。")
    except Exception as e:
        logger.critical(f"❌ 数据库连接失败，请检查 MySQL 是否运行及密码是否正确: {e}")
        exit(1)

init_database()

# ==================== 工具函数 ====================
def get_client_ip(req):
    return (req.headers.get('X-Forwarded-For') or
            req.headers.get('X-Real-IP') or
            req.remote_addr).split(',')[0].strip()

def safe_json_dumps(data):
    """安全地将数据转为 JSON 字符串，失败时返回 '{}'"""
    try:
        return json.dumps(data, ensure_ascii=False, separators=(',', ':'))
    except Exception as e:
        logger.warning(f"JSON dumps 失败: {e}")
        return '{}'

def safe_json_loads(text, default=dict):
    """安全地从字符串解析 JSON，失败时返回默认值"""
    if not text or not text.strip():
        return default() if callable(default) else default
    try:
        return json.loads(text)
    except Exception as e:
        logger.warning(f"JSON loads 失败: {e}")
        return default() if callable(default) else default

# ==================== 核心中间件 ====================
@app.before_request
def handle_request():
    if request.path.startswith(('/static', '/admin')):
        return None

    db = SessionLocal()
    try:
        method = request.method
        path = request.path

        # 查询 method + path 是否已存在
        existing = db.query(WebRequest).filter_by(method=method, path=path).first()

        if not existing:
            # ========== 新请求：插入记录 ==========
            try:
                json_data = request.get_json(silent=True)
            except:
                json_data = None

            # 过滤敏感头
            headers = {
                k: v for k, v in request.headers
                if k.lower() not in ['cookie', 'authorization']
            }

            new_record = WebRequest(
                method=method,
                path=path,
                query_string=request.query_string.decode('utf-8') if request.query_string else '',
                header=safe_json_dumps(headers),
                args=safe_json_dumps(dict(request.args)),
                form=safe_json_dumps(dict(request.form) if request.form else {}),
                json=safe_json_dumps(json_data),
                response=safe_json_dumps({"message": "welcome!"})
            )
            db.add(new_record)
            db.commit()
            logger.info(f"🆕 新请求记录: {method} {path} 来自 {get_client_ip(request)}")
            return jsonify("welcome!")

        else:
            # ========== 已存在请求：准备回放 response ==========
            resp_data = safe_json_loads(existing.response, default=dict)
            request._replay_response = resp_data
            request._db_session = db
            logger.info(f"🔁 回放响应: {method} {path}")
            return None  # 继续执行 after_request

    except exc.IntegrityError:
        db.rollback()
        logger.warning(f"⚠️ 并发请求，已存在: {method} {path}")
        return jsonify("welcome!")
    except Exception as e:
        db.rollback()
        logger.error(f"❌ 处理请求失败 {request.method} {request.path}: {e}")
        return jsonify({"error": "服务器内部错误"}), 500

@app.after_request
def replay_response(response):
    db = getattr(request, '_db_session', None)
    replay_data = getattr(request, '_replay_response', None)

    if replay_data is not None:
        try:
            resp = jsonify(replay_data)
            response.close()
            if db:
                db.close()
            return resp
        except Exception as e:
            logger.error(f"❌ 响应回放失败: {e}")
            if db:
                db.close()
            return jsonify({"error": "回放失败"}), 500
    elif db:
        db.close()
    return response

# ==================== 管理接口 ====================
@app.route('/admin/list', methods=['GET'])
def list_records():
    db = SessionLocal()
    try:
        records = db.query(WebRequest).all()
        data = []
        for r in records:
            data.append({
                "id": r.id,
                "method": r.method,
                "path": r.path,
                "query_string": r.query_string,
                "header": safe_json_loads(r.header, default=dict),
                "args": safe_json_loads(r.args, default=dict),
                "form": safe_json_loads(r.form, default=dict),
                "json": safe_json_loads(r.json, default=dict),
                "response": safe_json_loads(r.response, default=dict),
                "created_at": r.created_at.isoformat()
            })
        return jsonify(data)
    except Exception as e:
        logger.error(f"❌ 查询记录失败: {e}")
        return jsonify({"error": "查询失败"}), 500
    finally:
        db.close()

@app.route('/admin/update-response', methods=['POST'])
def update_response():
    db = SessionLocal()
    try:
        # 必须是 application/json
        if not request.is_json:
            return jsonify({"error": "Content-Type 必须是 application/json"}), 415

        data = request.get_json()
        method = data.get('method', 'GET').upper()
        path = data.get('path')
        new_response = data.get('response')

        if not path or new_response is None:
            return jsonify({"error": "缺少 path 或 response"}), 400

        record = db.query(WebRequest).filter_by(method=method, path=path).first()
        if record:
            # 序列化 response 并更新
            record.response = safe_json_dumps(new_response)
            db.commit()
            logger.info(f"🔧 更新响应: {method} {path}")
            return jsonify({"message": "更新成功"})
        else:
            return jsonify({"error": "路径未找到"}), 404
    except Exception as e:
        db.rollback()
        logger.error(f"❌ 更新响应失败: {e}")
        return jsonify({"error": "更新失败"}), 500
    finally:
        db.close()

# ==================== 示例接口 ====================
@app.route('/')
def index():
    return jsonify({"message": "Hello World"})

# ==================== 启动服务 ====================
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=SERVER_PORT, debug=False)